import cv2
import numpy as np
import os
import poimage


def check_file_exists(file_path):
    """检查文件是否存在（不是目录）"""
    return os.path.isfile(file_path)


# 支持jpg、png等所有图片格式
def process_image(input_image, output_image):
    poimage.del_watermark(
        input_image=input_image,
        output_image=output_image)


def read_image(image_path):
    try:
        image = cv2.imread(image_path)
        if image is None:
            raise ValueError(f"Failed to read image from {image_path}")
        return image
    except Exception as e:
        print(f"An error occurred while reading the image: {e}")


def detect_watermark(image):
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    edges = cv2.Canny(gray, 50, 150)
    contours, _ = cv2.findContours(edges, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    for contour in contours:
        x, y, w, h = cv2.boundingRect(contour)
        if w * h > 100:  # 假设水印区域大于100像素
            return x, y, w, h
    return None


def remove_watermark(image_path, output_path):
    if check_file_exists(image_path):
        print(f"文件 '{image_path}' 存在。")
    else:
        print(f"文件 '{image_path}' 不存在。")

    img = read_image(image_path)
    watermark_area = detect_watermark(img)
    if watermark_area:
        x, y, w, h = watermark_area
        mask = np.zeros(img.shape[:2], np.uint8)
        mask[y:y + h, x:x + w] = 255
        dst = cv2.inpaint(img, mask, 3, cv2.INPAINT_TELEA)
        cv2.imwrite(output_path, dst)


# 批量处理
image_paths = [r'13.jpg']
for i, path in enumerate(image_paths):
    process_image(path, f"{path}.jpg")
    # remove_watermark(path, f"{path}.jpg")
